#!/usr/bin/env python2.5
#coding: utf8

#Author: Gang Li
#Date: 2010-04-10 18:20

import socket
import sys
import select
import time
import cPickle as _p

from Queue import Queue as _q, Empty
from threading import Thread as _t
from threading import RLock
from StringIO import StringIO as _s

from core.utils import *

import logging
logging.basicConfig(format="%(name)s %(asctime)s %(levelname)s %(message)s", level=logging.INFO)
log = logging.getLogger("transport")
log.setLevel(logging.INFO)

class IncompleteCmd(Exception):
    def __init__(self, s):
        Exception.__init__(self, s)

class InvalidCmdData(Exception):
    def __init__(self, s):
        Exception.__init__(self, s)

def _safe_read(sock, amt):
    s = ''
    while amt > 0:
        chunk = sock.recv(amt)
        if not chunk:
            return None
        s = s + chunk
        amt = amt - len(chunk)
    return s

def write_cmd(sock, id, data=None):
    if data is None:
        data = ''

    out = _s()
    _p.dump(data, out)
    data = out.getvalue()
#    print data
#    data = data.encode("utf8")

    s = "%s:%s:%s" % (id, len(data), data)
    log.debug("send cmd: %s %s" % (s, sock))

    sock.send(s)

def read_cmd(sock):
    _cmd = None
    _len = None
    value = ''
    while 1:
        v = sock.recv(1)
        log.debug("read ... %s" % v)

        if not v:
            raise IncompleteCmd("read cmd error: %s-%s-%s-%s" % (_cmd, _len, value, v))
        
        if v == ':':
            if _cmd is None:
                _cmd = value
                value = ''
                continue
            
            if _len is None:
               _len = int(value)
               value = _safe_read(sock, _len)
               value = _p.loads(value)
               if value is None:
                   raise IncompleteCmd()

            log.debug("recv cmd: %s len: %s data: %s %s" % (_cmd, _len, value, sock))
            return _cmd, value
        else:
            value += v

_r, _w = read_cmd, write_cmd

class CmdDispatcher():
    filters = []
    def dispatch(self, msg):
        for f in self.filters:
            if f.accept_cmd(msg):
                f.on_cmd(msg)
                
    def add_filter(self, filter):
        self.filters.append(filter)

class CmdFilter():

    def accept_cmd(self, msg):
        """
            return True, means this cmd can be filter by another filter
            otherwise, this cmd must be stop propagation 
        """
        return True

    def on_cmd(self, msg):
        pass


class Server:
    """
    free
    alloc
    wait_task
    """

    dispatcher = CmdDispatcher()

    def __init__(self, binding_ip, port, handler=None):
        self.running = True
        self.port = port
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.bind((binding_ip, port))

        #tcp use this line 
        self.sock.listen(1)

        self.nodes = {}

        self.last_check = 0

        self.waits = [self.sock]
        self.writes = {}

        self.r_queue = _q()
        self.w_queue = _q()

        self.handler = handler

        _t(target=self._start_loop, args=[self]).start()
        _t(target=self._start_dispatch, args=[self]).start()

    def add_cmd_observer(self, handler):
        self.dispatcher.add_filter(handler)

    def _start_dispatch(self, kw):
        while self.running:
            try:
                msg = self.r_queue.get(timeout=2)
                self.dispatcher.dispatch(msg)
            except Empty:
                pass
            except BaseException, e:
                import traceback
                traceback.print_exc()
                log.warning("dispatch msg error: %s" , e)

    def _start_loop(self, kw):
        log.info("waiting for connections ... @ %s" , self.sock.getsockname())
        while self.running:
            try:
                self.process()
                self._write()
            except BaseException, msg:
                log.warning("service loop error: %s", msg)

    def process(self):
        n = select.select(self.waits, self.writes.keys(), [], 2)
        for skt in n[1]:
            try:
                funcs = self.writes[skt]
                chunk = funcs[0]()
                if chunk is None:
                    if len(funcs) > 1:funcs[1]()
                    del self.writes[skt]
                    continue

                skt.sendall(chunk)
            except BaseException, msg:
                log.warning("send data error: %s", msg)
                self.clean_client(skt)

        for skt in n[0]:
            if skt == self.sock:
                new_sock, (remote_ip, remote_port) = self.sock.accept()
#                i, p = new_sock.getsockname()
#                print "client socket", i, p
                log.debug("new connection from %s:%s" % (remote_ip, remote_port))
                self.nodes[new_sock] = {"last_seen":time.time()}
                self.waits.append(new_sock)

                self.handler and self.handler.connected(new_sock)
            else:
                try:
                    log.debug("read data from %s", skt.getsockname())
                    id, data = _r(skt)
                    #read cmd, add cmd to queue
                    self.r_queue.put((id, data, skt))
                    self.nodes[skt]["last_seen"] = time.time()
                except BaseException, msg:
                    log.warning("error: %s", msg)
                    self.clear_client(skt)

    def touch_write(self, skt, get_data, callback=None):
        if not get_data:raise Exception("get_data can not be None")
        self.writes[skt] = (get_data, callback)

    def write(self, skt, id, data=None): 
        self.w_queue.put((skt, id, data))

    def write_all(self, id, data=None): 
        for skt in self.nodes:
            self.w_queue.put((skt, id, data))

    def _write(self):
        if self.w_queue.empty():return
        try:
            skt, id, data = self.w_queue.get()
            _w(skt, id, data)
            self.nodes[skt]["last_seen"] = time.time()
        except BaseException, e:
            log.warning("write cmd error: %s" % e)

    def clear_client(self, skt):
        log.debug("clean client node")
        self.waits.remove(skt)
        del self.nodes[skt]
        del self.writes[skt]
        skt.close()
        try:
            self.handler and self.handler.disconnected(skt)
        except BaseException, e:
            log.warning("clear client error:%s", e)

    def close(self):
        self.running = False
        for k in self.nodes:
            self.clear_client(skt)

#    def verify_client(self, skt):
#        id, data = _r(skt)
#        if id == VERIFY_REQ:
#            _w(skt, VERIFY_RSP)
#        else:
#            print "error verify cmd", id, len, data, "expect", VERIFY_REQ
#            return
#
#        id, data = _r(skt)
#        if id == VERIFY_OK:
#            return True

class Client:
    dispatcher = CmdDispatcher()
    def __init__(self, srv_ip, srv_port, handler):
        try:
            self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.sock.connect((srv_ip, srv_port))
            self.last_check = time.time()
            self.last_seen = self.last_check
            self.r_queue = _q()
            self.w_queue = _q()
            
            self.handler = handler
            
            _t(target=self._start_loop, args=[self]).start()
            _t(target=self._start_dispatch, args=[self]).start()
        except socket.error, e:
            log.warning("socket error")

    def add_cmd_observer(self, handler):
        self.dispatcher.add_filter(handler)

#    def verify(self):
#        _w(self.sock, VERIFY_REQ)
#
#        id, data = _r(self.sock)
#        if id == VERIFY_RSP:
#            _w(self.sock, VERIFY_OK)
#            log.info("handshake success")
#            return True
#        else:
#            log.info("handshake failed")
#            return

    def write(self, id, data=None): 
        self.w_queue.put((id, data))

    def _write(self):
        if self.w_queue.empty():return
        try:
            _w(self.sock, *self.w_queue.get())
            self.last_seen = time.time()
        except BaseException, e:
            log.warning("write cmd error: %s" % e)

    def _start_loop(self, kw):
#        if not self.verify():
#            self.clear_connection()
#            return

        self.handler and self.handler.connected(self.sock)

        while self.sock:
            try:
                self._process()
                self._write()
            except BaseException, msg:
                log.warning(msg)

        log.info("process thd stop")
    def _start_dispatch(self, kw):
        while self.sock:
            try:
                msg = self.r_queue.get(timeout=2)
                self.dispatcher.dispatch(msg)
            except Empty:
                pass
            except BaseException, e:
                log.warning("dispatch msg error: %s" , e)
        log.info("dispatch thd stop")

    def clear_connection(self):
        log.info("server has gone!!!")
        #TODO retry to connect server
        self.sock.close()

        try:
            self.handler and self.handler.disconnected(self.sock)
        except BaseException, e :
            log.warning("clear connection %s", e)

        self.sock = None

    def _process(self):
        r, w, o = select.select([self.sock], [], [], 0.5)
        if r:
            for sock in r:
                try:
                    id, data = _r(sock)
                except BaseException, e:
                    log.warning("process error:%s", e)
                    self.clear_connection()
                    continue
                self.r_queue.put((id, data, sock))
                self.last_seen = time.time()

    def close(self):
        if self.sock:
            try:
                self.sock.close()
            except:
                pass
            self.sock = None

class HBEntry():
    def __init__(self, sender, on_failure):
        self.sender = sender
        self.on_failure = on_failure
        self.last_seen = time.time()
        self.last_check = time.time()

class HeartBeat(CmdFilter):
    def __init__(self):
        self.interval = 8
        self.seen_interval = 32
        self.check_interval = 24

        self.running = True
        self.observed = {}
        _t(target=self._do_hb, args=[self]).start()

    def add_observed(self, obj, sender, on_failure):
        self.observed[obj] = HBEntry(sender, on_failure)

    def _do_hb(self, kw):
        while self.running:
            time.sleep(self.interval)

            cur = time.time()
            for obj, obs in self.observed.items():
                self.__do_hb(obj, obs, cur)

    def __do_hb(self, obj, obs, cur):
        try:
            if cur - obs.last_check > self.check_interval:
                if cur - obs.last_seen > self.seen_interval * 3:
                    log.warning("hb failure, try to call 'on_failure'")
                    obs.on_failure and obs.on_failure(obj)
                    del self.observed[obj]
                elif cur - obs.last_seen > self.seen_interval:
                    obs.sender(HEART_BEAT_REQ)
                else:
                    pass

                obs.last_check = cur
        except BaseException, e:
            log.warning("check hb error: %s", e)

    def close(self):
        self.running = False

    def on_cmd(self, msg):
        id, data, skt = msg

        log.debug("heart beat monitor, recv cmd: %s", id)
        if id == HEART_BEAT_REQ:
            log.debug("hb rsp sent")
            self.observed[skt].sender(HEART_BEAT_RSP)
        elif id == HEART_BEAT_RSP:
            pass

        self.observed[skt].last_seen = time.time()

    def accept_cmd(self, msg):
        id = msg[0]
        return id in [HEART_BEAT_REQ, HEART_BEAT_RSP]
